# import torch
# import numpy as np
# from torch.utils.data import DataLoader
# from torch import nn
# from model.model import Transfollower, lstm_model, nn_model
# from data.dataset import get_data
# from config import Settings, HighDSettings
# import os
# import matplotlib.pyplot as plt
# import random  # 用于随机选择样本
#
# # 选择数据集和模型
# DATASET = 'highD'  # ['SH', 'NGSIM', 'highD']
#
# if DATASET == 'highD':
#     settings = HighDSettings()
# else:
#     settings = Settings()
#
# MODEL = 'transfollower'  # ['transfollower','lstm', 'nn']
#
# exp_name = f'{DATASET}_{MODEL}'
# save = f'checkpoints/{exp_name}_model.pt'
#
# # 参数
# SEQ_LEN = settings.SEQ_LEN
# LABEL_LEN = settings.LABEL_LEN
# PRED_LEN = settings.PRED_LEN
# BATCH_SIZE = settings.BATCH_SIZE
# lr = settings.lr
# T = settings.T  # 数据采样间隔
#
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#
# # 选择模型
# if MODEL == 'transfollower':
#     model = Transfollower(config=settings).to(device)
# elif MODEL == 'lstm':
#     model = lstm_model(config=settings).to(device)
# elif MODEL == 'nn':
#     model = nn_model(config=settings).to(device)
#
# # 载入整个模型
# model = torch.load(save)  # 去掉weights_only=True，加载完整模型
# model.to(device)  # 如果你使用的是 GPU，确保模型被正确移到 GPU
#
# # 设置评估模式
# model.eval()
#
# # 获取测试数据加载器
# _, _, test_loader = get_data(data_name='data', config=settings)
#
# # 损失函数
# criterion = nn.MSELoss()
#
#
# # 测试函数
# def test(data_loader):
#     total_loss = []
#     all_pred = []  # 用于存储所有预测结果
#     all_labels = []  # 用于存储所有真实标签
#
#     with torch.no_grad():
#         for i, item in enumerate(data_loader):
#             enc_inp = item['his'].float().to(device)  # (256,50,9)
#             batch_y = item['svSpd'].float()  # (256,149,1)
#             y_label = batch_y[:, -PRED_LEN:, :].to(device)  # (256,137,1)
#             batch_y_mark = item['pre'].float().to(device)  # (256,149,1)
#
#             dec_inp = torch.zeros([batch_y.shape[0], PRED_LEN, batch_y.shape[-1]]).float() + \
#                       batch_y[:, :LABEL_LEN, :].mean(axis=1, keepdim=True)  # 前一个的形状是[256,137,1]填充的后一个的形状是[256,1,1]
#             dec_inp = torch.cat([batch_y[:, :LABEL_LEN, :], dec_inp], dim=1).float().to(device)  # 拼接[256,149,1]
#             dec_inp = torch.cat([dec_inp, batch_y_mark], axis=-1)  # 拼接成[256,149,2]
#
#             # 模型推理
#             if MODEL == 'nn':
#                 out = model(dec_inp)
#             elif MODEL == 'transfollower':
#                 out = model(enc_inp, dec_inp)[0]  # 输出的第1个对应  model.py里的out[256,137,1]
#             else:
#                 out = model(enc_inp, dec_inp)
#
#             pred_lon_speed = out[:, :, 0].unsqueeze(-1)  # 纵向速度预测 (256, 137, 1)
#
#             # 纵向速度损失
#             lon_speed_loss = criterion(pred_lon_speed, y_label)
#
#             # 总损失
#             loss = lon_speed_loss
#
#             total_loss.append(loss.item())
#
#             # 存储预测和真实标签
#             all_pred.append(pred_lon_speed.cpu().numpy())
#             all_labels.append(y_label.cpu().numpy())
#
#     return np.mean(total_loss), np.concatenate(all_pred, axis=0), np.concatenate(all_labels, axis=0)
#
#
# # 可视化函数
# def visualize(predictions, labels):
#     # 随机选择一个样本
#     idx = random.randint(0, predictions.shape[0] - 1)
#
#     plt.figure(figsize=(10, 6))
#     plt.plot(labels[idx], label=f"True {idx}", color='blue')
#     plt.plot(predictions[idx], label=f"Predicted {idx}", color='red')
#
#     plt.xlabel('Time Step')
#     plt.ylabel('Velocity')
#     plt.title('True vs Predicted Longitudinal Speed')
#     plt.legend()
#     plt.grid(True)
#     plt.show()
#
#
# # 运行测试并可视化
# def main():
#     test_loss, predictions, labels = test(test_loader)
#     print(f"Test Loss: {test_loss:.7f}")
#
#     # 随机抽取一条并可视化
#     visualize(predictions, labels)
#
#
# if __name__ == '__main__':
#     main()
#################################平滑处理
# import torch
# import numpy as np
# from torch.utils.data import DataLoader
# from torch import nn
# from model.model import Transfollower, lstm_model, nn_model
# from data.dataset import get_data
# from config import Settings, HighDSettings
# import os
# import matplotlib.pyplot as plt
# import random
# from scipy.signal import savgol_filter
#
#
# # 数据异常值剔除函数（基于标准差）
# def remove_outliers(data, threshold=3):
#     mean = np.mean(data)
#     std_dev = np.std(data)
#     return data[np.abs(data - mean) <= threshold * std_dev]
#
# # 数据平滑函数（Savitzky-Golay滤波）
# def smooth_data_savgol(data, window_length=15, polyorder=3):
#     return savgol_filter(data, window_length=window_length, polyorder=polyorder, axis=0)
#
#
# # 选择数据集和模型
# DATASET = 'highD'  # ['SH', 'NGSIM', 'highD']
# if DATASET == 'highD':
#     settings = HighDSettings()
# else:
#     settings = Settings()
#
# MODEL = 'transfollower'  # ['transfollower','lstm', 'nn']
#
# exp_name = f'{DATASET}_{MODEL}'
# save = f'checkpoints/{exp_name}_model.pt'
#
# # 参数
# SEQ_LEN = settings.SEQ_LEN
# LABEL_LEN = settings.LABEL_LEN
# PRED_LEN = settings.PRED_LEN
# BATCH_SIZE = settings.BATCH_SIZE
# lr = settings.lr
# T = settings.T  # 数据采样间隔
#
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#
# # 选择模型
# if MODEL == 'transfollower':
#     model = Transfollower(config=settings).to(device)
# elif MODEL == 'lstm':
#     model = lstm_model(config=settings).to(device)
# elif MODEL == 'nn':
#     model = nn_model(config=settings).to(device)
#
# # 载入整个模型
# model = torch.load(save)  # 去掉weights_only=True，加载完整模型
# model.to(device)  # 如果你使用的是 GPU，确保模型被正确移到 GPU
#
# # 设置评估模式
# model.eval()
#
# # 获取测试数据加载器
# _, _, test_loader = get_data(data_name='data', config=settings)
#
# # 损失函数
# criterion = nn.MSELoss()
#
#
# # 测试函数
# def test(data_loader):
#     total_loss = []
#     all_pred = []  # 用于存储所有预测结果
#     all_labels = []  # 用于存储所有真实标签
#
#     with torch.no_grad():
#         for i, item in enumerate(data_loader):
#             enc_inp = item['his'].float().to(device)  # (256,50,9)
#             batch_y = item['svSpd'].float()  # (256,149,1)
#             y_label = batch_y[:, -PRED_LEN:, :].to(device)  # (256,137,1)
#             batch_y_mark = item['pre'].float().to(device)  # (256,149,1)
#
#             dec_inp = torch.zeros([batch_y.shape[0], PRED_LEN, batch_y.shape[-1]]).float() + \
#                       batch_y[:, :LABEL_LEN, :].mean(axis=1, keepdim=True)  # 前一个的形状是[256,137,1]填充的后一个的形状是[256,1,1]
#             dec_inp = torch.cat([batch_y[:, :LABEL_LEN, :], dec_inp], dim=1).float().to(device)  # 拼接[256,149,1]
#             dec_inp = torch.cat([dec_inp, batch_y_mark], axis=-1)  # 拼接成[256,149,2]
#
#             # 模型推理
#             if MODEL == 'nn':
#                 out = model(dec_inp)
#             elif MODEL == 'transfollower':
#                 out = model(enc_inp, dec_inp)[0]  # 输出的第1个对应  model.py里的out[256,137,1]
#             else:
#                 out = model(enc_inp, dec_inp)
#
#             pred_lon_speed = out[:, :, 0].unsqueeze(-1)  # 纵向速度预测 (256, 137, 1)
#
#             # 纵向速度损失
#             lon_speed_loss = criterion(pred_lon_speed, y_label)
#
#             # 总损失
#             loss = lon_speed_loss
#
#             total_loss.append(loss.item())
#
#             # 存储预测和真实标签
#             all_pred.append(pred_lon_speed.cpu().numpy())
#             all_labels.append(y_label.cpu().numpy())
#
#     return np.mean(total_loss), np.concatenate(all_pred, axis=0), np.concatenate(all_labels, axis=0)
#
#
# # 可视化函数
# def visualize(predictions, labels):
#     # 随机选择一个样本
#     idx = random.randint(0, predictions.shape[0] - 1)
#
#     # 获取样本的真实值和预测值
#     true_data = labels[idx]
#     pred_data = predictions[idx]
#
#     # 处理：去除异常值和进行平滑处理
#     true_data_cleaned = remove_outliers(true_data)
#     pred_data_cleaned = remove_outliers(pred_data)
#
#     # 平滑处理
#     true_data_smoothed = smooth_data_savgol(true_data_cleaned)
#     pred_data_smoothed = smooth_data_savgol(pred_data_cleaned)
#
#     # 绘制对比图
#     plt.figure(figsize=(10, 6))
#     plt.plot(true_data, label="True (Raw)", color='blue', alpha=0.6)
#     plt.plot(pred_data, label="Predicted (Raw)", color='red', alpha=0.6)
#
#     plt.plot(true_data_smoothed, label="True (Smoothed)", color='green', linestyle='--')
#     plt.plot(pred_data_smoothed, label="Predicted (Smoothed)", color='orange', linestyle='--')
#
#     plt.xlabel('Time Step')
#     plt.ylabel('Velocity')
#     plt.title('True vs Predicted Longitudinal Speed')
#     plt.legend()
#     plt.grid(True)
#     plt.show()
#
#
# # 运行测试并可视化
# def main():
#     test_loss, predictions, labels = test(test_loader)
#     print(f"Test Loss: {test_loss:.7f}")
#
#     # 随机抽取一条并可视化
#     visualize(predictions, labels)
#
#
# if __name__ == '__main__':
#     main()
import numpy as np
#######################全部图片保存
# import torch
# import numpy as np
# from torch.utils.data import DataLoader
# from torch import nn
# from model.model import Transfollower, lstm_model, nn_model
# from data.dataset import get_data
# from config import Settings, HighDSettings
# import os
# import matplotlib.pyplot as plt
# import random  # 用于随机选择样本
#
# # 选择数据集和模型
# DATASET = 'highD'  # ['SH', 'NGSIM', 'highD']
#
# if DATASET == 'highD':
#     settings = HighDSettings()
# else:
#     settings = Settings()
#
# MODEL = 'transfollower'  # ['transfollower','lstm', 'nn']
#
# exp_name = f'{DATASET}_{MODEL}'
# save = f'checkpoints/{exp_name}_model.pt'
#
# # 参数
# SEQ_LEN = settings.SEQ_LEN
# LABEL_LEN = settings.LABEL_LEN
# PRED_LEN = settings.PRED_LEN
# BATCH_SIZE = settings.BATCH_SIZE
# lr = settings.lr
# T = settings.T  # 数据采样间隔
#
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#
# # 选择模型
# if MODEL == 'transfollower':
#     model = Transfollower(config=settings).to(device)
# elif MODEL == 'lstm':
#     model = lstm_model(config=settings).to(device)
# elif MODEL == 'nn':
#     model = nn_model(config=settings).to(device)
#
# # 载入整个模型
# model = torch.load(save)  # 去掉weights_only=True，加载完整模型
# model.to(device)  # 如果你使用的是 GPU，确保模型被正确移到 GPU
#
# # 设置评估模式
# model.eval()
#
# # 获取测试数据加载器
# _, _, test_loader = get_data(data_name='data', config=settings)
#
# # 损失函数
# criterion = nn.MSELoss()
#
#
# # 测试函数
# def test(data_loader):
#     total_loss = []
#     all_pred = []  # 用于存储所有预测结果
#     all_labels = []  # 用于存储所有真实标签
#
#     with torch.no_grad():
#         for i, item in enumerate(data_loader):
#             enc_inp = item['his'].float().to(device)  # (256,50,9)
#             batch_y = item['svSpd'].float()  # (256,149,1)
#             y_label = batch_y[:, -PRED_LEN:, :].to(device)  # (256,137,1)
#             batch_y_mark = item['pre'].float().to(device)  # (256,149,1)
#
#             dec_inp = torch.zeros([batch_y.shape[0], PRED_LEN, batch_y.shape[-1]]).float() + \
#                       batch_y[:, :LABEL_LEN, :].mean(axis=1, keepdim=True)  # 前一个的形状是[256,137,1]填充的后一个的形状是[256,1,1]
#             dec_inp = torch.cat([batch_y[:, :LABEL_LEN, :], dec_inp], dim=1).float().to(device)  # 拼接[256,149,1]
#             dec_inp = torch.cat([dec_inp, batch_y_mark], axis=-1)  # 拼接成[256,149,2]
#
#             # 模型推理
#             if MODEL == 'nn':
#                 out = model(dec_inp)
#             elif MODEL == 'transfollower':
#                 out = model(enc_inp, dec_inp)[0]  # 输出的第1个对应  model.py里的out[256,137,1]
#             else:
#                 out = model(enc_inp, dec_inp)
#
#             pred_lon_speed = out[:, :, 0].unsqueeze(-1)  # 纵向速度预测 (256, 137, 1)
#
#             # 纵向速度损失
#             lon_speed_loss = criterion(pred_lon_speed, y_label)
#
#             # 总损失
#             loss = lon_speed_loss
#
#             total_loss.append(loss.item())
#
#             # 存储预测和真实标签
#             all_pred.append(pred_lon_speed.cpu().numpy())
#             all_labels.append(y_label.cpu().numpy())
#
#     return np.mean(total_loss), np.concatenate(all_pred, axis=0), np.concatenate(all_labels, axis=0)
#
#
# # 可视化函数
# def visualize_and_save(predictions, labels, save_dir):
#     # 确保保存目录存在
#     os.makedirs(save_dir, exist_ok=True)
#
#     for idx in range(predictions.shape[0]):
#         plt.figure(figsize=(10, 6))
#         plt.plot(labels[idx], label=f"True {idx}", color='blue')
#         plt.plot(predictions[idx], label=f"Predicted {idx}", color='red')
#
#         plt.xlabel('Time Step')
#         plt.ylabel('Velocity')
#         plt.title(f'True vs Predicted Longitudinal Speed (Sample {idx})')
#         plt.legend()
#         plt.grid(True)
#
#         # 保存图片到指定文件夹
#         plt.savefig(os.path.join(save_dir, f'sample_{idx}.png'))
#         plt.close()
#
#
# # 运行测试并保存所有图像
# def main():
#     test_loss, predictions, labels = test(test_loader)
#     print(f"Test Loss: {test_loss:.7f}")
#
#     # 设置保存图片的文件夹路径
#     save_dir = 'test_results2'
#
#     # 可视化并保存每个样本的对比图
#     visualize_and_save(predictions, labels, save_dir)
#
#
# if __name__ == '__main__':
#     main()
#######################带bev的测试代码##########################
import torch
from matplotlib import pyplot as plt
from torch import nn
from tensorboardX import SummaryWriter
from data.dataset import get_data
from model.model import *
from config import *
import os
import warnings
from train import *

warnings.filterwarnings("ignore", category=FutureWarning)

# 设置
DATASET = 'highD'
settings = HighDSettings() if DATASET == 'highD' else Settings()
exp_name = f'{DATASET}_transfollower'
save = f'checkpoints/{exp_name}_model.pt'
writer = SummaryWriter(f'runs/{exp_name}')

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Transfollower(config=settings).to(device)

mse_loss = nn.MSELoss()

# 加载训练好的模型
model = torch.load(save)
model.to(device)
model.eval()

# 获取测试数据加载器
_, _, test_loader = get_data(settings, data_name='data')


def construct_dec_inp(batch_x, pre_all, label_len, pred_len):
    dec_inp = torch.zeros([batch_x.shape[0], pred_len, batch_x.shape[-1]], device=batch_x.device) + \
              batch_x[:, :label_len, :].mean(dim=1, keepdim=True)
    dec_inp = torch.cat([batch_x[:, :label_len, :], dec_inp], dim=1)
    return torch.cat([dec_inp, pre_all], dim=-1)


# 测试函数
def test(data_loader):
    model.eval()
    total_loss = []
    all_pred = []  # 用于存储所有预测结果
    all_labels = []  # 用于存储所有真实标签

    with torch.no_grad():
        for item in data_loader:
            batch_x = item['latSpd'].to(device)
            x_label = batch_x[:, -settings.PRED_LEN:, :]
            roadbev_his = item['roadbev_his'].to(device)
            roadbev_pre = item['roadbev_pre'].to(device)
            data_pre = item['pre'].to(device)
            data_pre = construct_dec_inp(batch_x, data_pre, settings.LABEL_LEN, settings.PRED_LEN)

            data_his = item['his'].to(device)
            enc_inp = data_his
            dec_inp = data_pre
            enc_inp = enc_inp.float()
            dec_inp = dec_inp.float()
            x_label = x_label.float()
            out = model(enc_inp.float(), dec_inp.float(), roadbev_his, roadbev_pre)

            pred_lat_speed = out[0][:, :, 0].unsqueeze(-1)
            loss = mse_loss(pred_lat_speed, x_label)
            total_loss.append(loss.item())

            # 存储预测和真实标签
            all_pred.append(pred_lat_speed.cpu().numpy())
            all_labels.append(x_label.cpu().numpy())

    return np.mean(total_loss), np.concatenate(all_pred, axis=0), np.concatenate(all_labels, axis=0)


# 可视化函数
def visualize_and_save(predictions, labels, save_dir):
    # 确保保存目录存在
    os.makedirs(save_dir, exist_ok=True)

    for idx in range(predictions.shape[0]):
        plt.figure(figsize=(10, 6))
        plt.plot(labels[idx], label=f"True {idx}", color='blue')
        plt.plot(predictions[idx], label=f"Predicted {idx}", color='red')

        plt.xlabel('Time Step')
        plt.ylabel('Velocity')
        plt.title(f'True vs Predicted Longitudinal Speed (Sample {idx})')
        plt.legend()
        plt.grid(True)

        # 保存图片到指定文件夹
        plt.savefig(os.path.join(save_dir, f'sample_{idx}.png'))
        plt.close()


# 运行测试并保存所有图像
def main():
    test_loss, predictions, labels = test(test_loader)
    print(f"Test Loss: {test_loss:.7f}")

    # 设置保存图片的文件夹路径
    save_dir = 'test_results3'

    # 可视化并保存每个样本的对比图
    visualize_and_save(predictions, labels, save_dir)


if __name__ == '__main__':
    main()
